import pickle
import numpy as np
import torch
import os

from Config.Config import ACTNET200V13_PKL,TRAIN_PROPOSALS_PKL,TEST_PROPOSALS_PKL,VAL_PROPOSALS_PKL,TSNscore_FEATURE_PATH
from DataSet.Tools_TSNscore import load_TSNscore_from_file,extract_TSN_feature
from DataSet.Dataset_TSNscore import Dataset_TSNscore

class Dataset_TSNscore_Class(Dataset_TSNscore):

    def tempo_bilinear(self,feat):
        return np.dot(np.transpose(feat),feat)

    def _random_sample_traing_item(self,batch=1024, random=32, hard_neg=32, maxnump=32, thr_ious=[0.7,0.3]):
        x = np.zeros((batch,202*202))
        y = np.ones((batch))*-1

        target_pos = batch-hard_neg-random
        pos_count = 0
        neg_count = 0
        rand_count = 0
        count = 0

        for pos, vid in enumerate(self.vids):
            vitem = self.groundtruth[vid]
            if not os.path.exists(TSNscore_FEATURE_PATH+'/feat/{}.h5'.format(vid)) or vid not in self.proposals or len(self.proposals[vid]) == 0:
                continue

            feat = load_TSNscore_from_file(vid)
            num_frame = vitem['numf']
            duration = vitem['duration']
            proposal = np.asarray(self.proposals[vid])

            max_ious = proposal[:,:,2].max(axis=1)
            max_ious_lv = proposal[:,:,2].argmax(axis=1)

            pos_indexes = np.nonzero(max_ious>thr_ious[0])[0]
            neg_indexes = np.nonzero(max_ious<thr_ious[1])[0]

            PYRAMID = [10000]

            pos_pairs = []
            for i in range(len(pos_indexes)):
                ind = pos_indexes[i]
                lv = max_ious_lv[ind]
                pos_pairs += [proposal[ind,lv:lv+1,:]]
            np.random.shuffle(pos_pairs)
            for pair in pos_pairs[:min(target_pos,maxnump)]:
                if pos_count < target_pos and count < batch:
                    data = extract_TSN_feature(feat, pair, num_frame, duration, PYRAMID=PYRAMID,repeat=False)
                    if data is not None :
                        x[count, ...].flat = self.tempo_bilinear(data).flatten()
                        y[count, ...] = pair[0,3]
                        pos_count += 1
                        count += 1

            # sample random negative around around positive
            rand_pairs = []
            for i in range(len(pos_indexes)):
                ind = pos_indexes[i]
                lv = max_ious_lv[ind]
                rand_pairs += [proposal[ind,i:i+1,:] for i in range(6) if i!=lv]
            np.random.shuffle(rand_pairs)
            for pair in rand_pairs[:min(random,maxnump)]:
                if rand_count<random and count < batch:
                    data = extract_TSN_feature(feat, pair, num_frame, duration, PYRAMID=PYRAMID,repeat=False)
                    if data is not None :
                        x[count, ...].flat = self.tempo_bilinear(data).flatten()
                        y[count, ...] = pair[0,3]
                        rand_count += 1
                        count += 1

            # sample hard negatives
            neg_pairs = []
            for i in range(len(neg_indexes)):
                ind = neg_indexes[i]
                neg_pairs += [ proposal[ind,i:i+1,] for i in range(6) ]
            np.random.shuffle(neg_pairs)
            for pair in neg_pairs[:min(hard_neg,maxnump)]:
                if neg_count < hard_neg and count < batch:
                    data = extract_TSN_feature(feat, pair, num_frame, duration, PYRAMID=PYRAMID,repeat=False)
                    if data is not None :
                        x[count, ...].flat = self.tempo_bilinear(data).flatten()
                        y[count, ...] = 0
                        neg_count += 1
                        count += 1

            if count == batch:
                # 位移动一下
                self.vids = self.vids[pos:] + self.vids[:pos]
                break

        return x,y

    def nextbatch(self):
        return self._random_sample_traing_item()

    def _enum_every_test_item(self):

        # 一次出来一组6个和对应的VID
        for vid,proposals in self.proposals.items():

            vitem = self.groundtruth[vid]
            if not os.path.exists(TSNscore_FEATURE_PATH+'/feat/{}.h5'.format(vid)) or vid not in self.proposals or len(self.proposals[vid]) == 0:
                continue
            feat = load_TSNscore_from_file(vid)

            num_frame = vitem['numf']
            duration = vitem['duration']
            proposals = np.asarray(proposals)

            # proposals: 64*6 64个锚点每个锚点6个不同长度
            for i in range(proposals.shape[0]):
                current = np.copy(proposals[i,...])
                # current shape 6x4
                current[:,3] = 0
                n = current.shape[0]
                inputx = extract_TSN_feature(feat,current,num_frame,duration,PYRAMID=[10000]*n,SpliteData=True,repeat=False)
                if inputx is None: continue
                INPT = np.zeros((n,202*202))
                index = []
                for j in range(n):
                    if inputx[j] is None: continue
                    index.append(j)
                    INPT[j,...] = self.tempo_bilinear(inputx[j]).flatten()
                current = current[index,:]
                INPT = INPT[index,:]
                yield [vid,current,INPT]


    def enum_genfucntion(self):
        return self._enum_every_test_item()

def test1():
    dataset = Dataset_TSNscore_Class('validation', 'train')
    x,y = dataset.nextbatch()
    # x: (1024,40804) y: (1024,)

def test2():
    dataset = Dataset_TSNscore_Class('validation', 'test')
    gen = dataset.enum_genfucntion()
    vid,cur,inputx = next(gen)

if __name__=='__main__':
    test2()
